import numpy as np
from scipy.stats import f, ncf


def repeated_measures_contrast_sample_size(
        alpha: float,
        M: int,
        means: list,
        coe: list,
        scale: float,  # "Scale = sqrt(MSE)"，在对比中常记作 D
        sd_each_level: float,  # 每个水平的标准差（文中 often = 10）
        cor: float,
        power: float,
        step: float = 0.01,
        max_iter: int = 9999
):
    """
    基于线性组合(Contrast)的单因素重复测量 ANOVA 样本量估计。
    (复刻 Overall & Doyle 方法，以及文中例1-4在 SAS 中的实现)
    -------------------------------------------------------------------
    参数说明:
      alpha        : 显著性水平（如 0.05）
      M            : 水平数（如有 4 个时间点 -> M=4）
      means        : 各水平(时间点)上的均值 [μ1, μ2, μ3, μ4, ...]
      coe          : Contrast 系数 [c1, c2, c3, c4, ...], 通常满足 sum(c_i)=0 或 1
      scale        : 在文中称为 "Scale, D = sqrt(MSE)"
      sd_each_level: 每个水平的标准差。例如文中常用 10，而 D=4.4721
      cor          : 各水平间相关系数 (ρ)
      power        : 目标检验功效 (0~1), 如 0.90
      step         : 迭代步长 (可用 1.0, 0.5, 0.01等)，决定搜索精度
      max_iter     : 最大迭代次数（防止无限循环）

    计算逻辑:
      1) Contrast 值: C = Σ( c_i * μ_i )
      2) SDD         : sqrt(Σ(c_i^2))
      3) 效应量 es   : |C| / ( scale * sd_each_level * sqrt(1 - cor) )
         ※ 注意这里 scale (D=√MSE) 和 sd_each_level 要同时乘进分母。
      4) 对比只有1个自由度 => df1=1, df2=(n - 1)*(M - 1)
         非中心参数 ncp = n * (es^2)
      5) 功效 = 1 - F_{ncf}(f_crit, df1, df2, ncp)
         若功效>=目标，则返回样本量 n。

    返回:
      (n_found, pw)
       - n_found: 满足功效时的最小样本量 (int)
       - pw     : 该 n 下实际计算到的功效值
    """
    if len(means) != M:
        raise ValueError("means 的长度必须与 M 一致.")
    if len(coe) != M:
        raise ValueError("coe (对比系数) 的长度必须与 M 一致.")
    if not (0 < alpha < 1):
        raise ValueError("alpha 必须在 (0,1).")
    if not (0 < power < 1):
        raise ValueError("power 必须在 (0,1).")

    # 1) 计算 Contrast
    means_arr = np.array(means, dtype=float)
    coe_arr = np.array(coe, dtype=float)
    C = np.sum(coe_arr * means_arr)

    # 2) 计算 SDD = sqrt( sum(c_i^2) )
    SDD = np.sqrt(np.sum(coe_arr ** 2))

    # 3) 效应量 es
    es = abs(C) / (scale * sd_each_level * np.sqrt(1 - cor))

    # 4) 迭代搜索最小 n
    n = 2
    for _ in range(max_iter):
        df1 = 1
        df2 = (n - 1) * (M - 1)
        if df2 <= 0:
            n += step
            continue

        # 非中心参数 ncp
        ncp = n * (es ** 2)

        # 中心 F 分布下临界值 (1 - alpha) 分位数
        f_crit = f.ppf(1 - alpha, df1, df2)
        if np.isinf(f_crit) or np.isnan(f_crit):
            n += step
            continue

        # 功效 = 1 - F_{ncf}(f_crit; df1, df2, ncp)
        cdf_val = ncf.cdf(f_crit, df1, df2, ncp)
        pw = 1 - cdf_val

        if pw >= power:
            # 找到满足需求的样本量
            return int(np.ceil(n)), pw

        n += step

    # 若循环完都找不到，则返回 (None, None)
    return None, None


# ---------------------- 测试示例 ----------------------
if __name__ == "__main__":
    # 假设你要复现文中“例 1–4”，其中:
    #   - 4 个水平 -> M=4
    #   - Contrast系数 coe = [-3, -1, +1, +3]
    #   - means (有些参考资料中：Mean1=55, Mean2=56.5, Mean3=58, Mean4=59.5) -- 仅举例
    #   - "Standard deviation at each level" = 10
    #   - "Scale, D = sqrt(MSE)" = 4.4721 (文中示例)
    #   - 相关系数 cor = 0.7
    #   - alpha=0.05, power=0.90

    alpha_test = 0.05
    M_test = 4
    means_test = [55, 56.5, 58, 59.5]  # 假定这么写，也可改成你实际想用的均值
    coe_test = [-3, -1, 1, 3]  # 线性对比系数
    scale_test = 4.4721  # sqrt(MSE)
    sd_each_level_test = 10  # 每个水平的标准差
    cor_test = 0.7
    power_test = 0.90

    n_est, pw_est = repeated_measures_contrast_sample_size(
        alpha=alpha_test,
        M=M_test,
        means=means_test,
        coe=coe_test,
        scale=scale_test,  # 与 sd_each_level 区分
        sd_each_level=sd_each_level_test,
        cor=cor_test,
        power=power_test,
        step=0.01  # 迭代增量
    )

    print("=== Contrast 样本量估计示例 ===")
    print(f"Contrast, C = sum(coe * means) = {np.sum(np.array(coe_test) * np.array(means_test))}")
    print(f"Scale(D) = {scale_test}, sd_each_level = {sd_each_level_test}, cor = {cor_test}")
    print(f"目标功效 = {power_test * 100}%")
    print("-----------------------------------")
    print(f"计算得到样本量 n = {n_est}")
    print(f"对应功效估计 ≈ {pw_est:.4f}")